# Spectral GCN + Attention Recovery + LSTM
# This code trains and tests the GNN model for the COVID-19 infection prediction in Tokyo
# Author: Jiawei Xue, August 26, 2021
# Step 1: read and pack the traning and testing data
# Step 2: training epoch, training process, testing
# Step 3: build the model = spectral GCN + Attention Recovery + LSTM
# Step 4: main function
# Step 5: evaluation
# Step 6: visualization

import os
import csv
import json
import copy
import time
import random
import string
import argparse
import numpy as np
import pandas as pd
import geopandas as gpd
import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
import matplotlib.pyplot as plt
from matplotlib import pyplot as plt
import torch.nn.functional as F
from spectral_T3_GCN_memory_light import SpecGCN
from spectral_T3_GCN_memory_light import SpecGCN_LSTM
#torch.set_printoptions(precision=8)
#hyperparameter for the setting
X_day, Y_day = 21,7
#START_DATE, END_DATE = '20200414','20210207'
#START_DATE, END_DATE = '20200808','20210603'
START_DATE, END_DATE = '20200720','20210515'
WINDOW_SIZE = 7

#hyperparameter for the learning
DROPOUT, ALPHA = 0.50, 0.20
NUM_EPOCHS, BATCH_SIZE, LEARNING_RATE = 100, 8, 0.0001
HIDDEN_DIM_1, OUT_DIM_1, HIDDEN_DIM_2 = 6,4,2
infection_normalize_ratio = 100.0
web_search_normalize_ratio = 100.0
train_ratio = 0.7
validate_ratio = 0.1

Step 1: read and pack the training and testing data

#1.total period (mobility+text): 
#from 20200201 to 20210620: (29+31+30+31+30+31+31+30+31+30+31)+(31+28+31+30+31+20)\
#= 335 + 171 = 506;
#2.number of zones: 23;
#3.infection period:
#20200331 to 20210620: (1+30+31+30+31+31+30+31+30+31)+(31+28+31+30+31+20) = 276 + 171 = 447.

#1. Mobility: functions 1.2 to 1.7
#2. Text:  functions 1.8 to 1.14
#3. InfectionL:  functions 1.15 
#4. Preprocess:  functions 1.16 to 1.24
#5. Learn:  functions 1.25 to 1.26

#function 1.1
#get the central areas of Tokyo (e.g., the Special wards of Tokyo)
#return: a 23 zone shapefile
def read_tokyo_23():
    folder = "/data/HSEES/xue/xue_codes/disease_prediction_ml/gml_code/present_model_version10/tokyo_23" 
    file = "tokyo_23zones.shp"
    path = os.path.join(folder,file) 
    data = gpd.read_file(path)   
    return data

##################1.Mobility#####################
#function 1.2
#get the average of two days' mobility (infection) records
def mob_inf_average(data, key1, key2):
    new_record = dict()
    record1, record2 = data[key1], data[key2]
    for i in record1:
        if i in record2:
            new_record[i] = (record1[i]+record2[i])/2.0
    return new_record

#function 1.3
#get the average of multiple days' mobility (infection) records
def mob_inf_average_multiple(data, keyList):
    new_record = dict()
    num_day = len(keyList)
    for i in range(num_day):
        record = data[keyList[i]]
        for zone_id in record:
            if zone_id not in list(new_record.keys()):
                new_record[zone_id] = record[zone_id]
            else:
                new_record[zone_id] += record[zone_id]
    for new_record_key in new_record:
        new_record[new_record_key] = new_record[new_record_key]*1.0/num_day
    return new_record

#function 1.4
#generate the dateList: [20200101, 20200102, ..., 20211231]
def generate_dateList():
    yearList = ["2020","2021"]
    monthList = ["0"+str(i+1) for i in range(9)] + ["10","11","12"]
    dayList = ["0"+str(i+1) for i in range(9)] + [str(i) for i in range(10,32)]
    day_2020_num = [31,29,31,30,31,30,31,31,30,31,30,31]
    day_2021_num = [31,28,31,30,31,30,31,31,30,31,30,31]
    date_2020, date_2021 = list(), list()
    for i in range(12):
        for j in range(day_2020_num[i]):
              date_2020.append(yearList[0] + monthList[i] + dayList[j])
        for j in range(day_2021_num[i]):
              date_2021.append(yearList[1] + monthList[i] + dayList[j])
    date_2020_2021 = date_2020 + date_2021
    return date_2020_2021

#function 1.5
#smooth the mobility (infection) data using the neighborhood average
#under a given window size 
#dateList: [20200101, 20200102, ..., 20211231]
def mob_inf_smooth(data, window_size, dateList):
    data_copy = copy.copy(data)
    data_key_list = list(data_copy.keys())
    for data_key in data_key_list:
        left = int(max(dateList.index(data_key)-(window_size-1)/2, 0))
        right = int(min(dateList.index(data_key)+(window_size-1)/2, len(dateList)-1))
        potential_neighbor = dateList[left:right+1]
        neighbor_data_key = list(set(data_key_list).intersection(set(potential_neighbor)))
        data_average = mob_inf_average_multiple(data_copy, neighbor_data_key)
        data[data_key] =  data_average
    return data

#function 1.6
#set the mobility (infection) of one day as zero
def mob_inf_average_null(data, key1, key2):
    new_record = dict()
    record1, record2 = data[key1], data[key2]
    for i in record1:
        if i in record2:
            new_record[i] = 0
    return new_record

#function 1.7
#read the mobility data from "mobility_feature_20200201.json"...
#return: all_mobility:{"20200201":{('123','123'):12345,...},...}
#20200201 to 20210620: 506 days
def read_mobility_data(jcode23):
    all_mobility = dict()
    mobilityFilePath = "/data/HSEES/xue/xue_codes/disease_prediction_ml/gml_code/"+\
    "present_model_version10/mobility_20210804"
    mobilityNameList = os.listdir(mobilityFilePath)
    for i in range(len(mobilityNameList)):
        day_mobility = dict()
        file_name = mobilityNameList[i] 
        if "20" in file_name:
            day = (file_name.split("_")[2]).split(".")[0]  #get the day
            file_path = mobilityFilePath + '/' + file_name
            f = open(file_path,)
            df_file = json.load(f)   #read the mobility file
            f.close()
            for key in df_file:
                origin, dest = key.split("_")[0], key.split("_")[1]
                if origin in jcode23 and dest in jcode23:
                    if origin == dest:
                        day_mobility[(origin, dest)] = 0.0     #ignore the inner-zone flow
                    else:
                        day_mobility[(origin, dest)] = df_file[key] 
            all_mobility[day] = day_mobility
    #missing data
    all_mobility["20201128"] = mob_inf_average(all_mobility,"20201127","20201129")
    all_mobility["20210104"] = mob_inf_average(all_mobility, "20210103","20210105")
    return all_mobility

##################2.Text#####################
#function 1.8
#get the average of two days' infection records
def text_average(data, key1, key2):
    new_record = dict()
    record1, record2 = data[key1], data[key2]
    for i in record1:
        if i in record2:
            zone_record1, zone_record2 = record1[i], record2[i]
            new_zone_record = dict()
            for j in zone_record1:
                if j in zone_record2:
                    new_zone_record[j] = (zone_record1[j] + zone_record2[j])/2.0
            new_record[i] = new_zone_record
    return new_record

#function 1.9
#get the average of multiple days' text records
def text_average_multiple(data, keyList):
    new_record = dict()
    num_day = len(keyList)
    for i in range(num_day):
        record = data[keyList[i]]
        for zone_id in record:                           #zone_id
            if zone_id not in new_record:
                new_record[zone_id] = dict()   
            for j in record[zone_id]:                    #symptom
                if j not in new_record[zone_id]:
                    new_record[zone_id][j] = record[zone_id][j]
                else: 
                    new_record[zone_id][j] += record[zone_id][j]
    for zone_id in new_record:
        for j in new_record[zone_id]:
            new_record[zone_id][j] = new_record[zone_id][j]*1.0/num_day 
    return new_record

#function 1.10
#smooth the text data using the neighborhood average
#under a given window size 
def text_smooth(data, window_size, dateList):
    data_copy = copy.copy(data)
    data_key_list = list(data_copy.keys())
    for data_key in data_key_list:
        left = int(max(dateList.index(data_key)-(window_size-1)/2, 0))
        right = int(min(dateList.index(data_key)+(window_size-1)/2, len(dateList)-1))
        potential_neighbor = dateList[left:right+1]
        neighbor_data_key = list(set(data_key_list).intersection(set(potential_neighbor)))
        data_average = text_average_multiple(data_copy, neighbor_data_key)
        data[data_key] =  data_average
    return data

#function 1.11
#read the number of user points
def read_point_json():
    with open('user_point/mobility_user_point.json') as point1:
        user_point1 = json.load(point1)
    with open('user_point/mobility_user_point_20210812.json') as point2:
        user_point2 = json.load(point2)
    user_point_all = dict()
    for i in user_point1:
        user_point_all[i] = user_point1[i]
    for i in user_point2:
        user_point_all[i] = user_point2[i]
    user_point_all["20201128"] = user_point_all["20201127"]  #data missing
    user_point_all["20210104"] = user_point_all["20210103"]  #data missing
    return user_point_all

#function 1.12
#normalize the text search by the number of user points.
def normalize_text_user(all_text, user_point_all):
    for day in all_text:
        if day in user_point_all:
            num_user = user_point_all[day]["num_user"]
            all_text_day_new = dict()
            all_text_day = all_text[day]
            for zone in all_text_day:
                if zone not in all_text_day_new:
                    all_text_day_new[zone] = dict()
                for sym in all_text_day[zone]:
                    all_text_day_new[zone][sym] = all_text_day[zone][sym]*1.0/num_user
            all_text[day] = all_text_day_new
    return all_text
    
#function 1.13
#read the text data
#20200201 to 20210620: 506 days
#all_text = {"20200211":{"123":{"code":3,"fever":2,...},...},...}
def read_text_data(jcode23):
    all_text = dict()
    textFilePath = "/data/HSEES/xue/xue_codes/disease_prediction_ml/gml_code/"+\
    "present_model_version10/text_20210804"
    textNameList = os.listdir(textFilePath)
    for i in range(len(textNameList)):
        day_text = dict()
        file_name = textNameList[i]
        if "20" in file_name:
            day = (file_name.split("_")[2]).split(".")[0]
            file_path = textFilePath + "/" + file_name
            f = open(file_path,)
            df_file = json.load(f)   #read the mobility file
            f.close()
            new_dict = dict()
            for key in df_file:
                if key in jcode23:
                    new_dict[key] = {key1:df_file[key][key1]*1.0*web_search_normalize_ratio for key1 in df_file[key]}
                    #new_dict[key] = df_file[key]*WEB_SEARCH_RATIO
            all_text[day] = new_dict
    all_text["20201030"] = text_average(all_text, "20201029", "20201031") #data missing
    return all_text

#function 1.14
#perform the min-max normalization for the text data.
def min_max_text_data(all_text,jcode23):
    #calculate the min_max
    #region_key: sym: [min,max]
    text_list = list(['痛み', '頭痛', '咳', '下痢', 'ストレス', '不安', \
                    '腹痛', 'めまい', '吐き気', '嘔吐', '筋肉痛', '動悸', \
                    '副鼻腔炎', '発疹', 'くしゃみ', '倦怠感', '寒気', '脱水', \
                    '中咽頭', '関節痛', '不眠症', '睡眠障害', '鼻漏', '片頭痛', \
                    '多汗症', 'ほてり', '胸痛', '発汗', '無気力', '呼吸困難', \
                    '喘鳴', '目の痛み', '体の痛み', '無嗅覚症', '耳の痛み', \
                    '錯乱', '見当識障害', '胸の圧迫感', '鼻の乾燥', '耳感染症', \
                    '味覚消失', '上気道感染症', '眼感染症', '食欲減少'])
    region_sym_min_max = dict()
    for key in jcode23:                          #initialize
        region_sym_min_max[key] = dict()
        for sym in text_list:
            region_sym_min_max[key][sym] = [1000000,0]  #min, max
    for day in all_text:                            #update
        for key in jcode23:
            for sym in text_list:
                if sym in all_text[day][key]:
                    count = all_text[day][key][sym]
                    if count < region_sym_min_max[key][sym][0]:
                        region_sym_min_max[key][sym][0] = count
                    if count > region_sym_min_max[key][sym][1]:
                        region_sym_min_max[key][sym][1] = count
    #print ("region_sym_min_max",region_sym_min_max)
    for key in jcode23:              #normalize
        for sym in text_list:
            min_count,max_count=region_sym_min_max[key][sym][0],region_sym_min_max[key][sym][1]
            for day in all_text:
                if sym in all_text[day][key]:
                    if max_count-min_count == 0:
                        all_text[day][key][sym] = 1
                    else:
                        all_text[day][key][sym] = (all_text[day][key][sym]-min_count)*1.0/(max_count-min_count)
                        #print("all_text[day][key][sym]",all_text[day][key][sym])
    return all_text

##################3.Infection#####################
#function 1.15
#read the infection data
#20200331 to 20210620: (1+30+31+30+31+31+30+31+30+31)+(31+28+31+30+31+20) = 276 + 171 = 447.
#all_infection = {"20200201":{"123":1,"123":2}}
def read_infection_data(jcode23):
    all_infection = dict()
    infection_path = "/data/HSEES/xue/xue_codes/disease_prediction_ml/gml_code/"+\
    "present_model_version10/patient_20210725.json"
    f = open(infection_path,)
    df_file = json.load(f)   #read the mobility file
    f.close()
    for zone_id in df_file:
        for one_day in df_file[zone_id]:
            daySplit = one_day.split("/") 
            year, month, day = daySplit[0], daySplit[1], daySplit[2]
            if len(month) == 1:
                month = "0" + month
            if len(day) == 1:
                day = "0" + day
            new_date = year + month + day
            if str(zone_id[0:5]) in jcode23:
                if new_date not in all_infection:
                    all_infection[new_date] = {zone_id[0:5]:df_file[zone_id][one_day]*1.0/infection_normalize_ratio}
                else:
                    all_infection[new_date][zone_id[0:5]] = df_file[zone_id][one_day]*1.0/infection_normalize_ratio
    #missing
    date_list = [str(20200316+i) for i in range(15)]
    for date in date_list:
        all_infection[date] = mob_inf_average(all_infection,'20200401','20200401')
    all_infection['20200514'] = mob_inf_average(all_infection,'20200513','20200515')
    all_infection['20200519'] = mob_inf_average(all_infection,'20200518','20200520')
    all_infection['20200523'] = mob_inf_average(all_infection,'20200522','20200524')
    all_infection['20200530'] = mob_inf_average(all_infection,'20200529','20200601')
    all_infection['20200531'] = mob_inf_average(all_infection,'20200529','20200601')
    all_infection['20201231'] = mob_inf_average(all_infection,'20201230','20210101')
    all_infection['20210611'] = mob_inf_average(all_infection,'20210610','20210612')
    #outlier
    all_infection['20200331'] = mob_inf_average(all_infection,'20200401','20200401')
    all_infection['20200910'] = mob_inf_average(all_infection,'20200909','20200912')
    all_infection['20200911'] = mob_inf_average(all_infection,'20200909','20200912')
    all_infection['20200511'] = mob_inf_average(all_infection,'20200510','20200512')
    all_infection['20201208'] = mob_inf_average(all_infection,'20201207','20201209')
    all_infection['20210208'] = mob_inf_average(all_infection,'20210207','20210209')
    all_infection['20210214'] = mob_inf_average(all_infection,'20210213','20210215')
    #calculate the subtraction
    all_infection_subtraction = dict()
    all_infection_subtraction['20200331'] = all_infection['20200331']
    all_keys = list(all_infection.keys())
    all_keys.sort()
    for i in range(len(all_keys)-1):
        record = dict()
        for j in all_infection[all_keys[i+1]]:
            record[j] = all_infection[all_keys[i+1]][j] - all_infection[all_keys[i]][j]
        all_infection_subtraction[all_keys[i+1]] = record
    return all_infection_subtraction, all_infection

##################4.Preprocess#####################
#function 1.16
#ensemble the mobility, text, and infection.
#all_mobility = {"20200201":{('123','123'):12345,...},...}
#all_text = {"20200201":{"123":{"cold":3,"fever":2,...},...},...}
#all_infection = {"20200316":{"123":1,"123":2}}
#all_x_y = {"0":[[mobility_1,text_1, ..., mobility_x_day,text_x_day], [infection_1,...,infection_y_day],\
#[infection_1,...,infection_x_day]],0}
#x_days, y_days: use x_days to predict y_days
def ensemble(all_mobility, all_text, all_infection, x_days, y_days, all_day_list):
    all_x_y = dict()
    for j in range(len(all_day_list) - x_days - y_days + 1):
        x_sample, y_sample, x_sample_infection = list(), list(), list()                   
        #add the data from all_day_list[0+j] to all_day_list[x_days-1+j]
        for k in range(x_days):
            day = all_day_list[k + j]
            x_sample.append(all_mobility[day])
            x_sample.append(all_text[day])  
            x_sample_infection.append(all_infection[day])             #concatenate with the infection data                       
        #add the data from all_day_list[x_days+j] to all_day_list[x_days+y_day-1+j]
        for k in range(y_days):
            day = all_day_list[x_days + k + j]
            y_sample.append(all_infection[day]) 
        all_x_y[str(j)] = [x_sample, y_sample, x_sample_infection,j]                          
    return all_x_y

#function 1.17
#split the data by train/validate/test = train_ratio/validation_ratio/(1-train_ratio-validation_ratio)
def split_data(all_x_y, train_ratio, validation_ratio):
    all_x_y_key = list(all_x_y.keys())
    n = len(all_x_y_key)
    n_train, n_validate = round(n*train_ratio), round(n*validation_ratio)
    n_test = n-n_train-n_validate
    train_key = [all_x_y[str(i)] for i in range(n_train)]
    validate_key = [all_x_y[str(i+n_train)] for i in range(n_validate)]
    test_key = [all_x_y[str(i+n_train+n_validate)] for i in range(n_test)]
    return train_key, validate_key, test_key

##function 1.18
#the second data split method
#split the data by train/validate/test = train_ratio/validation_ratio/(1-train_ratio-validation_ratio)
def split_data_2(all_x_y, train_ratio, validation_ratio):
    all_x_y_key = list(all_x_y.keys())
    n = len(all_x_y_key)
    n_train, n_validate = round(n*train_ratio), round(n*validation_ratio)
    n_test = n-n_train-n_validate
    train_list, validate_list = list(), list()
    
    train_validate_key = [all_x_y[str(i)] for i in range(n_train+n_validate)]
    train_key, validate_key = list(), list()
    for i in range(len(train_validate_key)):
        if i % 9 == 8:
            validate_key.append(all_x_y[str(i)])
            validate_list.append(i)
        else:
            train_key.append(all_x_y[str(i)])
            train_list.append(i)
    test_key = [all_x_y[str(i+n_train+n_validate)] for i in range(n_test)]
    return train_key, validate_key, test_key, train_list, validate_list

##function 1.19
#the third data split method
#split the data by train/validate/test = train_ratio/validation_ratio/(1-train_ratio-validation_ratio)
def split_data_3(all_x_y, train_ratio, validation_ratio):
    all_x_y_key = list(all_x_y.keys())
    n = len(all_x_y_key)
    n_train, n_validate = round(n*train_ratio), round(n*validation_ratio)
    n_test = n - n_train  - n_validate
    train_list, validate_list = list(), list()
    
    train_validate_key = [all_x_y[str(i)] for i in range(n_train + n_validate)]
    train_key, validate_key = list(), list()
    for i in range(len(train_validate_key)):
        if (n_train + n_validate-i) % 2 == 0 and (n_train + n_validate-i) <= 2*n_validate:
            validate_key.append(all_x_y[str(i)])
            validate_list.append(i)
        else:
            train_key.append(all_x_y[str(i)])
            train_list.append(i)
    test_key = [all_x_y[str(i+n_train+n_validate)] for i in range(n_test)]
    return train_key, validate_key, test_key, train_list, validate_list

##function 1.20
#find the mobility data starting from the day, which is x_days before the start_date
#start_date = "20200331", x_days = 7
def sort_date(all_mobility, start_date, x_days): 
    mobility_date_list = list(all_mobility.keys())
    mobility_date_list.sort()
    idx = mobility_date_list.index(start_date)
    mobility_date_cut = mobility_date_list[idx-x_days:] 
    return mobility_date_cut

#function 1.21
#find the mobility data starting from the day, which is x_days before the start_date,
#ending at the day, which is y_days after the end_date
#start_date = "20200331", x_days = 7
def sort_date_2(all_mobility, start_date, x_days, end_date, y_days): 
    mobility_date_list = list(all_mobility.keys())
    mobility_date_list.sort()
    idx = mobility_date_list.index(start_date)
    idx2 = mobility_date_list.index(end_date)
    mobility_date_cut = mobility_date_list[idx-x_days:idx2+y_days] 
    return mobility_date_cut

#function 1.22
#get the mappings from zone id to id, text id to id.
#get zone_text_to_idx 
def get_zone_text_to_idx(all_infection):
    zone_list = list(set(all_infection["20200401"].keys()))
    text_list = list(['痛み', '頭痛', '咳', '下痢', 'ストレス', '不安', \
                    '腹痛', 'めまい'])
    zone_list.sort()
    zone_dict = {str(zone_list[i]):i for i in range(len(zone_list))}
    text_dict = {str(text_list[i]):i for i in range(len(text_list))}
    return zone_dict, text_dict

#function 1.23
#change the data format to matrix
#zoneid_to_idx = {"13101":0, "13102":1, ..., "13102":22}
#sym_to_idx = {"cough":0}
#mobility: {('13101', '13101'): 709973, ...}
#text: {'13101': {'痛み': 51,...},...}  text
#infection: {'13101': 50, '13102': 137, '13103': 401,...} 
#data_type = {"mobility", "text", "infection"}
def to_matrix(zoneid_to_idx, sym_to_idx, input_data, data_type):
    n_zone, n_text = len(zoneid_to_idx), len(sym_to_idx)
    if data_type == "mobility":
        result = np.zeros((n_zone, n_zone))
        for key in input_data:    
            from_id, to_id = key[0], key[1]
            from_idx, to_idx = zoneid_to_idx[from_id], zoneid_to_idx[to_id]
            result[from_idx][to_idx] += input_data[key]
    if data_type == "text":
        result = np.zeros((n_zone, n_text))
        for key1 in input_data:
            for key2 in input_data[key1]:
                if key1 in list(zoneid_to_idx.keys()) and key2 in list(sym_to_idx.keys()):
                    zone_idx, text_idx = zoneid_to_idx[key1], sym_to_idx[key2]
                    result[zone_idx][text_idx] += input_data[key1][key2]
    if data_type == "infection":
        result = np.zeros(n_zone)
        for key in input_data:
            zone_idx = zoneid_to_idx[key]
            result[zone_idx] += input_data[key]
    return result

#function 1.24
#change the data to the matrix format
def change_to_matrix(data, zoneid_to_idx, sym_to_idx):
    data_result = list()
    for i in range(len(data)):
        combine1, combine2 = list(), list()
        combine3 = list()                                    #NEW
        mobility_text = data[i][0]
        x_infection_all = data[i][2]          #the x_days infection data
        day_order =  data[i][3] #NEW          the order of the day
        for j in range(round(len(mobility_text)*1.0/2)):
            mobility, text = mobility_text[2*j], mobility_text[2*j+1]
            x_infection =  x_infection_all[j]   #NEW
            new_mobility = to_matrix(zoneid_to_idx, sym_to_idx, mobility, "mobility")
            new_text = to_matrix(zoneid_to_idx, sym_to_idx, text, "text")
            combine1.append(new_mobility)
            combine1.append(new_text) 
            new_x_infection = to_matrix(zoneid_to_idx, sym_to_idx, x_infection, "infection") #NEW
            combine3.append(new_x_infection)   #NEW
        for j in range(len(data[i][1])):
            infection = data[i][1][j]                                                          
            new_infection = to_matrix(zoneid_to_idx, sym_to_idx, infection, "infection")
            combine2.append(new_infection)                                               
        data_result.append([combine1,combine2,combine3,day_order])    #mobility/text; infection_y; infection_x; day_order
    return data_result   

##################5.learn#####################
#function 1.25
def visual_loss(e_losses, vali_loss, test_loss):
    plt.figure(figsize=(4,3), dpi=300)
    x = range(len(e_losses))
    y1,y2,y3 = copy.copy(e_losses), copy.copy(vali_loss), copy.copy(test_loss)
    plt.plot(x,y1,linewidth=1, label="train")
    plt.plot(x,y2,linewidth=1, label="validate")
    plt.plot(x,y3,linewidth=1, label="test")
    plt.legend()
    plt.title('Loss decline on entire training/validation/testing data')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    #plt.savefig('final_f6.png',bbox_inches = 'tight')
    plt.show()

#function 1.26
def visual_loss_train(e_losses):
    plt.figure(figsize=(4,3), dpi=300)
    x = range(len(e_losses))
    y1 = copy.copy(e_losses)
    plt.plot(x,y1,linewidth=1, label="train")
    plt.legend()
    plt.title('Loss decline on entire training data')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    #plt.savefig('final_f6.png',bbox_inches = 'tight')
    plt.show()

Step 2: training epoch, training process, testing

#function 2.1
#normalize each column of the input mobility matrix as one
def normalize_column_one(input_matrix):
    column_sum = np.sum(input_matrix, axis=0)
    row_num, column_num = len(input_matrix), len(input_matrix[0])
    for i in range(row_num):
        for j in range(column_num):
             input_matrix[i][j] = input_matrix[i][j]*1.0/column_sum[j]
    return input_matrix

#function 2.2
#evalute the trained_model on validation or testing data.
def validate_test_process(trained_model, vali_test_data):
    criterion = nn.MSELoss()
    vali_test_y = [vali_test_data[i][1] for i in range(len(vali_test_data))]
    y_real = torch.tensor(vali_test_y)
    
    vali_test_x = [vali_test_data[i] for i in range(len(vali_test_data))]
    vali_test_x = convertAdj(vali_test_x)
    y_hat = trained_model.run_specGCN_lstm(vali_test_x)                                          
    loss = criterion(y_hat.float(), y_real.float())            ###Calculate the loss  
    return loss, y_hat, y_real 

#function 2.3
#convert the mobility matrix in x_batch in a following way
#normalize the flow between zones so that the in-flow of each zone is 1.
def convertAdj(x_batch):
    #x_batch：(n_batch, 0/1, 2*i+1)
    x_batch_new = copy.copy(x_batch)
    n_batch = len(x_batch)
    days = round(len(x_batch[0][0])/2)
    for i in range(n_batch):
        for j in range(days):
            mobility_matrix = x_batch[i][0][2*j]
            x_batch_new[i][0][2*j] = normalize_column_one(mobility_matrix)   #20210818
    return x_batch_new

#function 2.4
#a training epoch
def train_epoch_option(model, opt, criterion, trainX_c, trainY_c, batch_size):  
    model.train()
    losses = []
    batch_num = 0
    for beg_i in range(0, len(trainX_c), batch_size):
        batch_num += 1
        if batch_num % 16 ==0:
            print ("batch_num: ", batch_num, "total batch number: ", int(len(trainX_c)/batch_size))
        x_batch = trainX_c[beg_i:beg_i+batch_size]        
        y_batch = torch.tensor(trainY_c[beg_i:beg_i+batch_size])   
        opt.zero_grad()
        x_batch = convertAdj(x_batch)   #conduct the column normalization
        y_hat = model.run_specGCN_lstm(x_batch)                          ###Attention
        loss = criterion(y_hat.float(), y_batch.float()) #MSE loss
        #opt.zero_grad()
        loss.backward()
        opt.step()
        losses.append(loss.data.numpy())
    return sum(losses)/float(len(losses)), model

#function 2.5
#multiple training epoch
def train_process(train_data, lr, num_epochs, net, criterion, bs, vali_data, test_data):
    opt = optim.Adam(net.parameters(), lr, betas = (0.9,0.999), weight_decay=0) 
    train_y = [train_data[i][1] for i in range(len(train_data))]
    e_losses = list()
    e_losses_vali = list()
    e_losses_test = list()
    time00 = time.time()
    for e in range(num_epochs):
        time1 = time.time()
        print ("current epoch: ",e, "total epoch: ", num_epochs)
        number_list = list(range(len(train_data)))       
        random.shuffle(number_list)
        trainX_sample = [train_data[number_list[j]] for j in range(len(number_list))]
        trainY_sample = [train_y[number_list[j]] for j in range(len(number_list))]
        loss, net =  train_epoch_option(net, opt, criterion, trainX_sample, trainY_sample, bs)  
        print ("train loss", loss*infection_normalize_ratio*infection_normalize_ratio)
        e_losses.append(loss*infection_normalize_ratio*infection_normalize_ratio)
        
        loss_vali, y_hat_vali, y_real_vali = validate_test_process(net, vali_data) 
        loss_test, y_hat_test, y_real_test = validate_test_process(net, test_data)
        e_losses_vali.append(float(loss_vali)*infection_normalize_ratio*infection_normalize_ratio)
        e_losses_test.append(float(loss_test)*infection_normalize_ratio*infection_normalize_ratio)
        
        print ("validate loss", float(loss_vali)*infection_normalize_ratio*infection_normalize_ratio)
        print ("test loss", float(loss_test)*infection_normalize_ratio*infection_normalize_ratio)
        if e>=2 and (e+1)%10 ==0:
            visual_loss(e_losses, e_losses_vali, e_losses_test)     
            visual_loss_train(e_losses) 
        time2 = time.time()
        print ("running time for this epoch:", time2 - time1)
        time01 = time.time()
        print ("---------------------------------------------------------------")
        print ("---------------------------------------------------------------")
        #print ("total running time until now:", time01 - time00)
        #print ("------------------------------------------------")
        #print("specGCN_weight", net.specGCN.layer1.W)
        #print("specGCN_weight_grad", net.specGCN.layer1.W.grad)
        #print ("------------------------------------------------")
        #print("memory decay matrix", net.v)
        #print("memory decay matrix grad", net.v.grad)
        #print ("------------------------------------------------")
        #print ("lstm weight", net.lstm.all_weights[0][0])
        #print ("lstm weight grad", net.lstm.all_weights[0][0].grad)
        #print ("------------------------------------------------")
        #print ("fc1.weight", net.fc1.weight)
        #print ("fc1 weight grd", net.fc1.weight.grad)
        #print ("---------------------------------------------------------------")
        #print ("---------------------------------------------------------------")
    return e_losses, net

Step 3: models

#function 3.1
def read_data():
    jcode23 = list(read_tokyo_23()["JCODE"])                    #1.1 get the tokyo 23 zone shapefile
    all_mobility = read_mobility_data(jcode23)                  #1.2 read the mobility data
    all_text = read_text_data(jcode23)                          #1.3 read the text data
    all_infection, all_infection_cum = read_infection_data(jcode23)                #1.4 read the infection data
    
    #smooth the data using 7-days average
    window_size = WINDOW_SIZE                 #20210818
    dateList = generate_dateList()  #20210818
    all_mobility = mob_inf_smooth(all_mobility, window_size, dateList) #20210818
    all_infection = mob_inf_smooth(all_infection, window_size, dateList)  #20210818
    
    #smooth, user, min-max.
    point_json = read_point_json()                           #20210821
    all_text = normalize_text_user(all_text, point_json)       #20210821
    all_text = text_smooth(all_text, window_size, dateList) #20210818
    all_text = min_max_text_data(all_text,jcode23)                 #20210820
    
    x_days, y_days =  X_day, Y_day
    mobility_date_cut = sort_date_2(all_mobility, START_DATE, x_days, END_DATE, y_days)
    all_x_y = ensemble(all_mobility, all_text, all_infection, x_days, y_days, mobility_date_cut)
    train_original, validate_original, test_original, train_list, validation_list =\
    split_data_3(all_x_y,train_ratio,validate_ratio)  
    zone_dict, text_dict = get_zone_text_to_idx(all_infection)                       #get zone_idx, text_idx
    train_x_y = change_to_matrix(train_original, zone_dict, text_dict)                   #get train
    print ("train_x_y_shape",len(train_x_y),"train_x_y_shape[0]",len(train_x_y[0]))
    validate_x_y = change_to_matrix(validate_original, zone_dict, text_dict)             #get validate
    test_x_y = change_to_matrix(test_original, zone_dict, text_dict)                     #get test
    
    print (len(train_x_y))  #300
    print (len(train_x_y[0][0])) #14
    print (np.shape(train_x_y[0][0][0])) #(23,23)
    print (np.shape(train_x_y[0][0][1])) #(23,43)
    #print ("---------------------------------finish data reading and preprocessing------------------------------------")
    return train_x_y, validate_x_y, test_x_y, all_mobility, all_infection, train_original, validate_original, test_original, train_list, validation_list

#function 3.2
#train the model
def model_train(train_x_y, vali_data, test_data):
    #3.2.1 define the model
    input_dim_1, hidden_dim_1, out_dim_1, hidden_dim_2 = len(train_x_y[0][0][1][1]),\
    HIDDEN_DIM_1, OUT_DIM_1, HIDDEN_DIM_2 
    dropout_1, alpha_1, N = DROPOUT, ALPHA, len(train_x_y[0][0][1])
    G_L_Model = SpecGCN_LSTM(X_day, Y_day, input_dim_1, hidden_dim_1, out_dim_1, hidden_dim_2, dropout_1,N)         ###Attention
    #3.2.2 train the model
    num_epochs, batch_size, learning_rate = NUM_EPOCHS, BATCH_SIZE, LEARNING_RATE                                                 #model train
    criterion = nn.MSELoss() 
    e_losses, trained_model = train_process(train_x_y, learning_rate, num_epochs, G_L_Model, criterion, batch_size,\
                          vali_data, test_data)
    return e_losses, trained_model

#function 3.3
#evaluate the error on validation (or testing) data.
def validate_test_process(trained_model, vali_test_data):
    criterion = nn.MSELoss()
    vali_test_y = [vali_test_data[i][1] for i in range(len(vali_test_data))]
    y_real = torch.tensor(vali_test_y)
    vali_test_x = [vali_test_data[i] for i in range(len(vali_test_data))]
    vali_test_x = convertAdj(vali_test_x)
    y_hat = trained_model.run_specGCN_lstm(vali_test_x)                                  ###Attention              
    loss = criterion(y_hat.float(), y_real.float())
    return loss, y_hat, y_real 

Step 4. Model implementation

#4.1
#read the data
train_x_y, validate_x_y, test_x_y, all_mobility, all_infection, \
train_original, validate_original, test_original, train_list, validation_list =\
read_data()
#train_x_y, validate_x_y, test_x_y = normalize(train_x_y, validate_x_y, test_x_y)

#train_x_y = train_x_y[0:30]
print (len(train_x_y))
print ("---------------------------------finish data preparation------------------------------------")

train_x_y_shape 210 train_x_y_shape[0] 4
210
42
(23, 23)
(23, 8)
210
---------------------------------finish data preparation------------------------------------

#4.2
#train the model
e_losses, trained_model = model_train(train_x_y, validate_x_y, test_x_y)
print ("---------------------------finish model training-------------------------")

current epoch:  0 total epoch:  100
batch_num:  16 total batch number:  26
train loss 357.03401219237736
validate loss 252.24311277270317
test loss 449.4777321815491
running time for this epoch: 8.320138692855835
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  1 total epoch:  100
batch_num:  16 total batch number:  26
train loss 355.98878369287206
validate loss 227.68227383494377
test loss 440.91474264860153
running time for this epoch: 6.923511743545532
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  2 total epoch:  100
batch_num:  16 total batch number:  26
train loss 349.74736954878875
validate loss 202.12285220623016
test loss 428.5801202058792
running time for this epoch: 6.99811863899231
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  3 total epoch:  100
batch_num:  16 total batch number:  26
train loss 330.0012192792363
validate loss 179.51944842934608
test loss 413.3238270878792
running time for this epoch: 6.999978065490723
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  4 total epoch:  100
batch_num:  16 total batch number:  26
train loss 302.4154137475071
validate loss 161.4270731806755
test loss 396.36190980672836
running time for this epoch: 6.956462144851685
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  5 total epoch:  100
batch_num:  16 total batch number:  26
train loss 289.93628367229746
validate loss 149.40989203751087
test loss 381.5268352627754
running time for this epoch: 6.960854768753052
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  6 total epoch:  100
batch_num:  16 total batch number:  26
train loss 278.326786916565
validate loss 138.13842087984085
test loss 363.25886845588684
running time for this epoch: 7.06347393989563
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  7 total epoch:  100
batch_num:  16 total batch number:  26
train loss 266.0494884131131
validate loss 130.89505955576897
test loss 344.991534948349
running time for this epoch: 7.018553972244263
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  8 total epoch:  100
batch_num:  16 total batch number:  26
train loss 254.39065087724617
validate loss 127.07180343568325
test loss 326.1954337358475
running time for this epoch: 7.021049737930298
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  9 total epoch:  100
batch_num:  16 total batch number:  26
train loss 242.83469861580267
validate loss 126.62786059081554
test loss 308.17894265055656

running time for this epoch: 7.406564235687256
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  10 total epoch:  100
batch_num:  16 total batch number:  26
train loss 256.9699728930438
validate loss 129.28261421620846
test loss 289.8408845067024
running time for this epoch: 7.028077125549316
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  11 total epoch:  100
batch_num:  16 total batch number:  26
train loss 221.99559077206587
validate loss 135.73058880865574
test loss 269.4999799132347
running time for this epoch: 7.066872596740723
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  12 total epoch:  100
batch_num:  16 total batch number:  26
train loss 217.35022120453692
validate loss 143.99699866771698
test loss 252.5603026151657
running time for this epoch: 7.106000900268555
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  13 total epoch:  100
batch_num:  16 total batch number:  26
train loss 207.5524029701396
validate loss 152.68538147211075
test loss 238.20698261260986
running time for this epoch: 7.131714344024658
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  14 total epoch:  100
batch_num:  16 total batch number:  26
train loss 199.55064170062542
validate loss 164.72997143864632
test loss 222.41419181227684
running time for this epoch: 7.1510984897613525
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  15 total epoch:  100
batch_num:  16 total batch number:  26
train loss 190.85997905306243
validate loss 176.96833238005638
test loss 208.66597071290016
running time for this epoch: 7.19052267074585
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  16 total epoch:  100
batch_num:  16 total batch number:  26
train loss 204.24130131249078
validate loss 188.2505603134632
test loss 197.0439963042736
running time for this epoch: 7.145964622497559
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  17 total epoch:  100
batch_num:  16 total batch number:  26
train loss 178.77194978710676
validate loss 206.29173144698143
test loss 182.44819715619087
running time for this epoch: 7.137323617935181
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  18 total epoch:  100
batch_num:  16 total batch number:  26
train loss 181.1765568951766
validate loss 221.71633318066597
test loss 170.95835879445076
running time for this epoch: 7.174534797668457
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  19 total epoch:  100
batch_num:  16 total batch number:  26
train loss 170.85973604547758
validate loss 230.4093912243843
test loss 163.46512362360954

running time for this epoch: 7.4162373542785645
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  20 total epoch:  100
batch_num:  16 total batch number:  26
train loss 169.37629526687994
validate loss 241.04353040456772
test loss 155.51865100860596
running time for this epoch: 7.154966354370117
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  21 total epoch:  100
batch_num:  16 total batch number:  26
train loss 176.6323746630439
validate loss 255.75341656804085
test loss 146.3465765118599
running time for this epoch: 7.24921178817749
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  22 total epoch:  100
batch_num:  16 total batch number:  26
train loss 160.82889952317433
validate loss 257.9483203589916
test loss 142.04062521457672
running time for this epoch: 7.124257564544678
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  23 total epoch:  100
batch_num:  16 total batch number:  26
train loss 170.49097507778143
validate loss 268.8448317348957
test loss 134.8944567143917
running time for this epoch: 7.080856800079346
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  24 total epoch:  100
batch_num:  16 total batch number:  26
train loss 169.82105301900043
validate loss 281.44383803009987
test loss 127.53787450492382
running time for this epoch: 7.139843463897705
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  25 total epoch:  100
batch_num:  16 total batch number:  26
train loss 152.87567033535905
validate loss 293.5783378779888
test loss 120.78985571861267
running time for this epoch: 7.056027412414551
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  26 total epoch:  100
batch_num:  16 total batch number:  26
train loss 150.4876538766203
validate loss 301.93693935871124
test loss 115.33712968230247
running time for this epoch: 7.076523542404175
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  27 total epoch:  100
batch_num:  16 total batch number:  26
train loss 147.593026480603
validate loss 306.01605772972107
test loss 111.45544238388538
running time for this epoch: 7.123220682144165
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  28 total epoch:  100
batch_num:  16 total batch number:  26
train loss 145.93748099826
validate loss 310.7818216085434
test loss 107.29316622018814
running time for this epoch: 7.070485591888428
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  29 total epoch:  100
batch_num:  16 total batch number:  26
train loss 165.3217469100599
validate loss 317.1512484550476
test loss 103.02319191396236

running time for this epoch: 7.395423889160156
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  30 total epoch:  100
batch_num:  16 total batch number:  26
train loss 144.07079528879237
validate loss 304.55516651272774
test loss 103.06932032108307
running time for this epoch: 7.148285150527954
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  31 total epoch:  100
batch_num:  16 total batch number:  26
train loss 140.40848542073812
validate loss 307.6302260160446
test loss 99.51038286089897
running time for this epoch: 7.119216203689575
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  32 total epoch:  100
batch_num:  16 total batch number:  26
train loss 141.41684467041935
validate loss 313.24736773967743
test loss 95.4945757985115
running time for this epoch: 7.149806499481201
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  33 total epoch:  100
batch_num:  16 total batch number:  26
train loss 137.10487927137703
validate loss 315.52717089653015
test loss 92.42561645805836
running time for this epoch: 7.228296995162964
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  34 total epoch:  100
batch_num:  16 total batch number:  26
train loss 135.1214934968286
validate loss 318.84409487247467
test loss 88.9424979686737
running time for this epoch: 7.135137557983398
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  35 total epoch:  100
batch_num:  16 total batch number:  26
train loss 141.17203428651447
validate loss 319.19848173856735
test loss 86.26560680568218
running time for this epoch: 7.133162021636963
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  36 total epoch:  100
batch_num:  16 total batch number:  26
train loss 134.51856500641614
validate loss 312.73525208234787
test loss 85.02469398081303
running time for this epoch: 7.196616888046265
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  37 total epoch:  100
batch_num:  16 total batch number:  26
train loss 144.90791993146692
validate loss 318.45003366470337
test loss 81.23816922307014
running time for this epoch: 7.121097087860107
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  38 total epoch:  100
batch_num:  16 total batch number:  26
train loss 128.72464661658913
validate loss 320.0293704867363
test loss 78.28203029930592
running time for this epoch: 7.161484479904175
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  39 total epoch:  100
batch_num:  16 total batch number:  26
train loss 136.98980496782397
validate loss 316.0464018583298
test loss 76.45037025213242

running time for this epoch: 7.480090141296387
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  40 total epoch:  100
batch_num:  16 total batch number:  26
train loss 124.87592011759128
validate loss 321.47467136383057
test loss 72.84836377948523
running time for this epoch: 7.0973124504089355
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  41 total epoch:  100
batch_num:  16 total batch number:  26
train loss 123.56040283554681
validate loss 323.4526887536049
test loss 70.20327262580395
running time for this epoch: 7.155606985092163
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  42 total epoch:  100
batch_num:  16 total batch number:  26
train loss 138.46566072768636
validate loss 322.35678285360336
test loss 68.0612213909626
running time for this epoch: 7.203847885131836
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  43 total epoch:  100
batch_num:  16 total batch number:  26
train loss 134.6487547615888
validate loss 309.07848849892616
test loss 67.82311480492353
running time for this epoch: 7.1560142040252686
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  44 total epoch:  100
batch_num:  16 total batch number:  26
train loss 119.31036504123497
validate loss 310.82984060049057
test loss 65.14216307550669
running time for this epoch: 7.149839401245117
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  45 total epoch:  100
batch_num:  16 total batch number:  26
train loss 117.6308683568129
validate loss 313.35726380348206
test loss 62.40152753889561
running time for this epoch: 7.144679069519043
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  46 total epoch:  100
batch_num:  16 total batch number:  26
train loss 116.4756065958904
validate loss 308.5041791200638
test loss 60.938610695302486
running time for this epoch: 7.122967004776001
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  47 total epoch:  100
batch_num:  16 total batch number:  26
train loss 114.72454846457198
validate loss 302.45570465922356
test loss 59.609925374388695
running time for this epoch: 7.1513447761535645
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  48 total epoch:  100
batch_num:  16 total batch number:  26
train loss 113.78509009739868
validate loss 301.866140216589
test loss 57.53634497523308
running time for this epoch: 7.245504379272461
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  49 total epoch:  100
batch_num:  16 total batch number:  26
train loss 111.8675625400135
validate loss 299.8986653983593
test loss 55.706300772726536

running time for this epoch: 7.52165675163269
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  50 total epoch:  100
batch_num:  16 total batch number:  26
train loss 110.97722617840324
validate loss 300.8941002190113
test loss 53.43499127775431
running time for this epoch: 7.172056674957275
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  51 total epoch:  100
batch_num:  16 total batch number:  26
train loss 107.75272355035499
validate loss 289.59888964891434
test loss 53.052064031362534
running time for this epoch: 7.175610303878784
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  52 total epoch:  100
batch_num:  16 total batch number:  26
train loss 120.04194643210482
validate loss 290.22589325904846
test loss 51.01550370454788
running time for this epoch: 7.10328221321106
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  53 total epoch:  100
batch_num:  16 total batch number:  26
train loss 113.22384450101741
validate loss 278.1008370220661
test loss 50.67579913884401
running time for this epoch: 7.14003324508667
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  54 total epoch:  100
batch_num:  16 total batch number:  26
train loss 103.17740522118079
validate loss 275.7668122649193
test loss 49.049523659050465
running time for this epoch: 7.200954437255859
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  55 total epoch:  100
batch_num:  16 total batch number:  26
train loss 107.2846754902491
validate loss 272.134430706501
test loss 47.605098225176334
running time for this epoch: 7.128945589065552
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  56 total epoch:  100
batch_num:  16 total batch number:  26
train loss 100.22725608130847
validate loss 273.1626480817795
test loss 45.677805319428444
running time for this epoch: 7.129566431045532
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  57 total epoch:  100
batch_num:  16 total batch number:  26
train loss 98.122448436226
validate loss 270.99432423710823
test loss 44.24599464982748
running time for this epoch: 7.185617446899414
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  58 total epoch:  100
batch_num:  16 total batch number:  26
train loss 96.77762976261202
validate loss 266.52783155441284
test loss 43.16109698265791
running time for this epoch: 7.156966686248779
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  59 total epoch:  100
batch_num:  16 total batch number:  26
train loss 95.16629687924352
validate loss 262.223519384861
test loss 42.062108404934406

running time for this epoch: 7.490160942077637
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  60 total epoch:  100
batch_num:  16 total batch number:  26
train loss 94.11117958801763
validate loss 258.3744190633297
test loss 40.97089171409607
running time for this epoch: 7.215563058853149
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  61 total epoch:  100
batch_num:  16 total batch number:  26
train loss 93.09004612818913
validate loss 251.89075618982315
test loss 40.234411135315895
running time for this epoch: 7.157025337219238
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  62 total epoch:  100
batch_num:  16 total batch number:  26
train loss 90.90685018303769
validate loss 252.24819779396057
test loss 38.71576162055135
running time for this epoch: 7.2094244956970215
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  63 total epoch:  100
batch_num:  16 total batch number:  26
train loss 93.61100598686824
validate loss 246.11791595816612
test loss 38.02052699029446
running time for this epoch: 7.238545656204224
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  64 total epoch:  100
batch_num:  16 total batch number:  26
train loss 90.59089723181117
validate loss 241.83107540011406
test loss 37.11158176884055
running time for this epoch: 7.190441608428955
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  65 total epoch:  100
batch_num:  16 total batch number:  26
train loss 87.27590613826007
validate loss 235.55541411042213
test loss 36.448086611926556
running time for this epoch: 7.1940319538116455
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  66 total epoch:  100
batch_num:  16 total batch number:  26
train loss 85.7277559461417
validate loss 231.23795166611671
test loss 35.64608283340931
running time for this epoch: 7.254073619842529
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  67 total epoch:  100
batch_num:  16 total batch number:  26
train loss 84.99051612387929
validate loss 226.44450888037682
test loss 34.928382374346256
running time for this epoch: 7.167032480239868
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  68 total epoch:  100
batch_num:  16 total batch number:  26
train loss 84.93126283779188
validate loss 223.2484705746174
test loss 34.036540891975164
running time for this epoch: 7.183419466018677
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  69 total epoch:  100
batch_num:  16 total batch number:  26
train loss 82.12181683457285
validate loss 221.33177146315575
test loss 33.07268023490906

running time for this epoch: 7.522334575653076
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  70 total epoch:  100
batch_num:  16 total batch number:  26
train loss 81.33021200558653
validate loss 219.0321497619152
test loss 32.218832056969404
running time for this epoch: 7.151828289031982
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  71 total epoch:  100
batch_num:  16 total batch number:  26
train loss 80.28885261241899
validate loss 208.48769694566727
test loss 32.24581712856889
running time for this epoch: 7.226071119308472
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  72 total epoch:  100
batch_num:  16 total batch number:  26
train loss 79.41873429377598
validate loss 208.59850570559502
test loss 31.14827675744891
running time for this epoch: 7.239787578582764
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  73 total epoch:  100
batch_num:  16 total batch number:  26
train loss 79.11618488530317
validate loss 200.75347274541855
test loss 30.979078728705645
running time for this epoch: 7.1704323291778564
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  74 total epoch:  100
batch_num:  16 total batch number:  26
train loss 77.19183539005893
validate loss 198.04613664746284
test loss 30.23384604603052
running time for this epoch: 7.194357395172119
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  75 total epoch:  100
batch_num:  16 total batch number:  26
train loss 77.99622184527969
validate loss 197.161715477705
test loss 29.3297809548676
running time for this epoch: 7.213475704193115
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  76 total epoch:  100
batch_num:  16 total batch number:  26
train loss 77.32575720486541
validate loss 193.6745084822178
test loss 28.857411816716194
running time for this epoch: 7.19715690612793
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  77 total epoch:  100
batch_num:  16 total batch number:  26
train loss 81.76107563216377
validate loss 187.7833902835846
test loss 28.56052713468671
running time for this epoch: 7.169971942901611
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  78 total epoch:  100
batch_num:  16 total batch number:  26
train loss 73.82699422089865
validate loss 188.79223614931107
test loss 27.541937306523323
running time for this epoch: 7.219801664352417
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  79 total epoch:  100
batch_num:  16 total batch number:  26
train loss 72.60068665750119
validate loss 180.74313178658485
test loss 27.583655901253223

running time for this epoch: 7.470415115356445
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  80 total epoch:  100
batch_num:  16 total batch number:  26
train loss 75.11080415160568
validate loss 179.38684672117233
test loss 26.952647604048252
running time for this epoch: 7.220332860946655
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  81 total epoch:  100
batch_num:  16 total batch number:  26
train loss 71.88682871249814
validate loss 173.31672832369804
test loss 26.871231384575367
running time for this epoch: 7.257111549377441
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  82 total epoch:  100
batch_num:  16 total batch number:  26
train loss 70.87179571735086
validate loss 169.83689740300179
test loss 26.499673258513212
running time for this epoch: 7.188196897506714
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  83 total epoch:  100
batch_num:  16 total batch number:  26
train loss 71.62857904202407
validate loss 169.96631398797035
test loss 25.773521047085524
running time for this epoch: 7.2106099128723145
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  84 total epoch:  100
batch_num:  16 total batch number:  26
train loss 71.75027162129818
validate loss 163.17971050739288
test loss 25.853856932371855
running time for this epoch: 7.18751335144043
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  85 total epoch:  100
batch_num:  16 total batch number:  26
train loss 68.84971042853539
validate loss 156.72048553824425
test loss 25.901703629642725
running time for this epoch: 7.180413722991943
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  86 total epoch:  100
batch_num:  16 total batch number:  26
train loss 67.95055159643569
validate loss 157.13050961494446
test loss 25.177968200296164
running time for this epoch: 7.223779916763306
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  87 total epoch:  100
batch_num:  16 total batch number:  26
train loss 67.17441894579679
validate loss 151.73123218119144
test loss 25.17478773370385
running time for this epoch: 7.25140643119812
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  88 total epoch:  100
batch_num:  16 total batch number:  26
train loss 66.74337295677375
validate loss 147.08087779581547
test loss 25.18519526347518
running time for this epoch: 7.201742649078369
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  89 total epoch:  100
batch_num:  16 total batch number:  26
train loss 80.32328625968486
validate loss 149.03774484992027
test loss 24.320031516253948

running time for this epoch: 7.5138890743255615
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  90 total epoch:  100
batch_num:  16 total batch number:  26
train loss 65.95647352730373
validate loss 147.0062229782343
test loss 23.895336780697107
running time for this epoch: 7.216394662857056
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  91 total epoch:  100
batch_num:  16 total batch number:  26
train loss 64.99853192104233
validate loss 145.33805660903454
test loss 23.562212008982897
running time for this epoch: 7.179095506668091
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  92 total epoch:  100
batch_num:  16 total batch number:  26
train loss 66.13109400927054
validate loss 141.44080691039562
test loss 23.536800872534513
running time for this epoch: 7.184220552444458
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  93 total epoch:  100
batch_num:  16 total batch number:  26
train loss 64.71746497684056
validate loss 138.5325100272894
test loss 23.357803001999855
running time for this epoch: 7.227690935134888
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  94 total epoch:  100
batch_num:  16 total batch number:  26
train loss 63.420148605170354
validate loss 136.5200523287058
test loss 23.111917544156313
running time for this epoch: 7.149551153182983
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  95 total epoch:  100
batch_num:  16 total batch number:  26
train loss 63.225304302587006
validate loss 134.51687060296535
test loss 22.886726073920727
running time for this epoch: 7.170605421066284
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  96 total epoch:  100
batch_num:  16 total batch number:  26
train loss 63.74724519749482
validate loss 131.67272321879864
test loss 22.752389777451754
running time for this epoch: 7.259721994400024
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  97 total epoch:  100
batch_num:  16 total batch number:  26
train loss 62.36158871885251
validate loss 126.8845982849598
test loss 22.919229231774807
running time for this epoch: 7.202257394790649
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  98 total epoch:  100
batch_num:  16 total batch number:  26
train loss 62.121781067819235
validate loss 125.59408321976662
test loss 22.628698498010635
running time for this epoch: 7.23844313621521
---------------------------------------------------------------
---------------------------------------------------------------
current epoch:  99 total epoch:  100
batch_num:  16 total batch number:  26
train loss 61.36823618025692
validate loss 125.34094974398613
test loss 22.254674695432186

running time for this epoch: 7.565415382385254
---------------------------------------------------------------
---------------------------------------------------------------
---------------------------finish model training-------------------------

#4.3 
print (len(train_x_y))
print (len(validate_x_y))
print (len(test_x_y))
#4.3.1 model validation
validation_result, validate_hat, validate_real = validate_test_process(trained_model, validate_x_y)
print ("---------------------------------finish model validation------------------------------------")
print (len(validate_hat))
print (len(validate_real))
#4.3.2 model testing
#4.4. model test
test_result, test_hat, test_real = validate_test_process(trained_model, test_x_y)
print ("---------------------------------finish model testing------------------------------------")
print (len(test_real))
print (len(test_hat))

210
30
60
---------------------------------finish model validation------------------------------------
30
30
---------------------------------finish model testing------------------------------------
60
60

Step 5: Evaluation

#5.1 RMSE, MAPE, MAE, RMSLE
def RMSELoss(yhat,y):
    return float(torch.sqrt(torch.mean((yhat-y)**2)))

def MAPELoss(yhat,y):
    return float(torch.mean(torch.div(torch.abs(yhat-y), y)))

def MAELoss(yhat,y):
    return float(torch.mean(torch.div(torch.abs(yhat-y), 1)))

def RMSLELoss(yhat,y):
    log_yhat = torch.log(yhat+1)
    log_y = torch.log(y+1)
    return float(torch.sqrt(torch.mean((log_yhat-log_y)**2)))    

#compute RMSE
rmse_validate = list()
rmse_test = list()
for i in range(len(validate_x_y)):
    rmse_validate.append(float(RMSELoss(validate_hat[i],validate_real[i])))
for i in range(len(test_x_y)):
    rmse_test.append(float(RMSELoss(test_hat[i],test_real[i])))
print ("rmse_validate mean", np.mean(rmse_validate))
print ("rmse_test mean", np.mean(rmse_test))

#compute MAE
mae_validate = list()
mae_test = list()
for i in range(len(validate_x_y)):
    mae_validate.append(float(MAELoss(validate_hat[i],validate_real[i])))
for i in range(len(test_x_y)):
    mae_test.append(float(MAELoss(test_hat[i],test_real[i])))
    
print ("mae_validate mean", np.mean(mae_validate))
print ("mae_test mean", np.mean(mae_test))

#show RMSE and MAE together
mae_validate, rmse_validate, mae_test, rmse_test =\
np.array(mae_validate)*infection_normalize_ratio, np.array(rmse_validate)*infection_normalize_ratio,\
np.array(mae_test)*infection_normalize_ratio, np.array(rmse_test)*infection_normalize_ratio
print ("-----------------------------------------")
print ("mae_validate mean", round(np.mean(mae_validate),3), "     rmse_validate mean", round(np.mean(rmse_validate),3))
print ("mae_test mean", round(np.mean(mae_test),3), "         rmse_test mean", round(np.mean(rmse_test),3))
print ("-----------------------------------------")

rmse_validate mean 0.09805288668460443
rmse_test mean 0.04455104110733852
mae_validate mean 0.07649717128826529
mae_test mean 0.033239021740926515
-----------------------------------------
mae_validate mean 7.65      rmse_validate mean 9.805
mae_test mean 3.324          rmse_test mean 4.455
-----------------------------------------

print(validate_hat[0][Y_day-1])
print(torch.sum(validate_hat[0][Y_day-1]))
print(validate_real[0][Y_day-1])
print(torch.sum(validate_real[0][Y_day-1]))

tensor([0.0877, 0.2061, 0.4705, 0.6335, 0.2285, 0.3244, 0.2688, 0.4920, 0.3762,
        0.4593, 0.8452, 1.0859, 0.4340, 0.5067, 0.6441, 0.3415, 0.3745, 0.2756,
        0.5924, 0.6864, 0.7495, 0.3367, 0.6467], grad_fn=<SelectBackward>)
tensor(11.0663, grad_fn=<SumBackward0>)
tensor([0.0443, 0.1186, 0.2243, 0.4386, 0.1600, 0.2257, 0.1629, 0.3100, 0.2729,
        0.1886, 0.6786, 0.7429, 0.2000, 0.3614, 0.5071, 0.2143, 0.2700, 0.1914,
        0.5457, 0.4757, 0.5086, 0.4914, 0.4657], dtype=torch.float64)
tensor(7.7986, dtype=torch.float64)

x = range(len(rmse_validate))
plt.figure(figsize=(8,2),dpi=300)
l1 = plt.plot(x, np.array(rmse_validate), 'ro-',linewidth=0.8, markersize=1.2, label='RMSE')
l2 = plt.plot(x, np.array(mae_validate), 'go-',linewidth=0.8, markersize=1.2, label='MAE')
plt.xlabel('Date from the first day of validation',fontsize=12)
plt.ylabel("RMSE/MAE daily new cases",fontsize=10)
my_y_ticks = np.arange(0,2100, 500)
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)
plt.legend()
plt.grid()
plt.show()

x = range(len(mae_test))
plt.figure(figsize=(8,2),dpi=300)
l1 = plt.plot(x, np.array(rmse_test), 'ro-',linewidth=0.8, markersize=1.2, label='RMSE')
l2 = plt.plot(x, np.array(mae_test), 'go-',linewidth=0.5, markersize=1.2, label='MAE')
plt.xlabel('Date from the first day of test',fontsize=12)
plt.ylabel("RMSE/MAE Daily new cases",fontsize=10)
my_y_ticks = np.arange(0,2100, 500)
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)
plt.legend()
plt.grid()
plt.show()

Correlation

from scipy import stats
#validate
y_days = Y_day
validate_hat_sum = [float(torch.sum(validate_hat[i][y_days-1])) for i in range(len(validate_hat))]
validate_real_sum = [float(torch.sum(validate_real[i][y_days-1])) for i in range(len(validate_real))]
print ("the correlation between validation: ", stats.pearsonr(validate_hat_sum, validate_real_sum)[0])
#test
test_hat_sum = [float(torch.sum(test_hat[i][y_days-1])) for i in range(len(test_hat))]
test_real_sum = [float(torch.sum(test_real[i][y_days-1])) for i in range(len(test_real))]
print ("the correlation between test: ", stats.pearsonr(test_hat_sum, test_real_sum)[0])
#train
train_result, train_hat, train_real = validate_test_process(trained_model, train_x_y)
train_hat_sum = [float(torch.sum(train_hat[i][0])) for i in range(len(train_hat))]
train_real_sum = [float(torch.sum(train_real[i][0])) for i in range(len(train_real))]
print ("the correlation between train: ", stats.pearsonr(train_hat_sum, train_real_sum)[0])

the correlation between validation:  0.9371378089459015
the correlation between test:  0.8407822153076807
the correlation between train:  0.9256750690241817

step 6. Visualization

y1List = [np.sum(list(train_original[i+1][1][Y_day-1].values())) for i in range(len(train_original)-1)]
y2List = [np.sum(list(validate_original[i][1][Y_day-1].values())) for i in range(len(validate_original))]
y2List_hat = [float(torch.sum(validate_hat[i][Y_day-1])) for i in range(len(validate_hat))]
y3List = [np.sum(list(test_original[i][1][Y_day-1].values())) for i in range(len(test_original))]
y3List_hat = [float(torch.sum(test_hat[i][Y_day-1])) for i in range(len(test_hat))]

#x1 = np.array(range(len(y1List)))
#x2 = np.array([len(y1List)+j for j in range(len(y2List))])
x1 = train_list
x2 = validation_list
x3 = np.array([len(y1List)+len(y2List)+j for j in range(len(y3List))])

plt.figure(figsize=(8,2),dpi=300)
l1 = plt.plot(x1[0: len(y1List)], np.array(y1List)*infection_normalize_ratio, 'ro-',linewidth=0.8, markersize=2.0, label='train')
l2 = plt.plot(x2, np.array(y2List)*infection_normalize_ratio, 'go-',linewidth=0.8, markersize=2.0, label='validate')
l3 = plt.plot(x2, np.array(y2List_hat)*infection_normalize_ratio, 'g-',linewidth=2, markersize=0.1, label='validate_predict')
l4 = plt.plot(x3, np.array(y3List)*infection_normalize_ratio, 'bo-',linewidth=0.8, markersize=2, label='test')
l5 = plt.plot(x3, np.array(y3List_hat)*infection_normalize_ratio, 'b-',linewidth=2, markersize=0.1, label='test_predict')
#plt.xlabel('Date from the first day of 2020/4/1',fontsize=12)
plt.ylabel("Daily infection cases",fontsize=10)
my_y_ticks = np.arange(0,2100, 500)
my_x_ticks = list()
summary = 0 
my_x_ticks.append(summary) 
for i in range(5):
    summary += 60
    my_x_ticks.append(summary) 
plt.xticks(my_x_ticks)
plt.yticks(my_y_ticks) 
plt.xticks(fontsize=8)
plt.yticks(fontsize=12)
plt.title("SpectralGCN")
plt.legend()
plt.grid()
#plt.savefig('sg_peak4_21_21_1feature_0005.pdf',bbox_inches = 'tight')
plt.show()

step 7: Regional visualization

def getPredictionPlot(k):
    #location k
    x_k = [i for i in range(len(test_real))]
    real_k = [test_real[i][y_days-1][k] for i in range(len(test_real))]
    predict_k = [test_hat[i][y_days-1][k] for i in range(len(test_hat))]
    plt.figure(figsize=(4,2.5), dpi=300)
    l1 = plt.plot(x_k, np.array(real_k)*infection_normalize_ratio, 'ro-',linewidth=0.8, markersize=2.0, label='real',alpha = 0.8)
    l2 = plt.plot(x_k, np.array(predict_k)*infection_normalize_ratio, 'o-',color='black',linewidth=0.8, markersize=2.0, alpha = 0.8, label='predict')
    #plt.xlabel('Date from the first day of 2020/4/1',fontsize=12)
    #plt.ylabel("Daily infection cases",fontsize=10)
    my_y_ticks = np.arange(0,100,40)
    my_x_ticks = list()
    summary = 0 
    my_x_ticks.append(summary) 
    for i in range(3):
        summary += 10
        my_x_ticks.append(summary) 
    plt.xticks(my_x_ticks)
    plt.yticks(my_y_ticks) 
    plt.xticks(fontsize = 14)
    plt.yticks(fontsize = 14)
    plt.title("Real and predict daily infection for region "+str(k))
    plt.legend()
    plt.grid()
    #plt.savefig('sg_peak4_21_21_1feature_0005.pdf',bbox_inches = 'tight')
    plt.show()

for i in range(23):
    getPredictionPlot(i)